import json
import os
import shutil
 
 
def labelme2yolo_seg(class_name, json_dir, labels_dir):
    """
        此函数用来将labelme软件标注好的json格式转换为yolov_seg中使用的txt格式
        :param json_dir: labelme标注好的*.json文件所在文件夹
        :param labels_dir: 转换好后的*.txt保存文件夹
        :param class_name: 数据集中的类别标签
        :return:
    """
    list_labels = []  # 存放json文件的列表
 
    # 0.创建保存转换结果的文件夹
    os.makedirs(labels_dir, exist_ok=False)
 
    # 1.获取目录下所有的labelme标注好的Json文件，存入列表中
    for files in os.listdir(json_dir):  # 遍历json文件夹下的所有json文件
        file = os.path.join(json_dir, files)  # 获取一个json文件
        list_labels.append(file)  # 将json文件名加入到列表中
 
    for labels in list_labels:  # 遍历所有json文件
        with open(labels, "r") as f:
            file_in = json.load(f)
            shapes = file_in["shapes"]
            print(labels)
 
        txt_filename = os.path.basename(labels).replace(".json", ".txt")
        txt_path = os.path.join(labels_dir, txt_filename)  # 使用labels_dir变量指定保存路径
 
        with open(txt_path, "w+") as file_handle:
            for shape in shapes:
                line_content = []  # 初始化一个空列表来存储每个形状的坐标信息
                line_content.append(str(class_name.index(shape['label'])))  # 添加类别索引
                # 添加坐标信息
                for point in shape["points"]:
                    x = point[0] / file_in["imageWidth"]
                    y = point[1] / file_in["imageHeight"]
                    line_content.append(str(x))
                    line_content.append(str(y))
                # 使用空格连接列表中的所有元素，并写入文件
                file_handle.write(" ".join(line_content) + "\n")

 
def labelme2yolo_det(class_name, json_dir, labels_dir):
    """
        此函数用来将labelme软件标注好的json格式转换为yolo_detect中使用的txt格式
        :param json_dir: labelme标注好的*.json文件所在文件夹
        :param labels_dir: 转换好后的*.txt保存文件夹
        :param class_name: 数据集中的类别标签数组
        :return:
    """
    list_labels = []  # 存放json文件的列表
    # 0.创建保存转换结果的文件夹
    os.makedirs(labels_dir, exist_ok=True)
 
    # 1.获取目录下所有的labelme标注好的Json文件，存入列表中
    for files in os.listdir(json_dir):  # 遍历json文件夹下的所有json文件
        file = os.path.join(json_dir, files)  # 获取一个json文件
        list_labels.append(file)  # 将json文件名加入到列表中
 
    for labels in list_labels:  # 遍历所有json文件
        with open(labels, "r",encoding='utf-8') as f:
            file_in = json.load(f)
            shapes = file_in["shapes"]

        if shapes == []:
            print("无对象：",labels)
            continue

        txt_filename = os.path.basename(labels).replace(".json", ".txt")
        txt_path = os.path.join(labels_dir, txt_filename)  # 使用labels_dir变量指定保存路径
 
        with open(txt_path, "w+") as file_handle:
            for shape in shapes:
                line_content = []  # 初始化一个空列表来存储每个形状的坐标信息
                line_content.append(str(class_name.index(shape['label'])))  # 添加类别索引
                [[x1, y1], [x2, y2]] = shape['points']
                x1, x2 = x1 / file_in['imageWidth'], x2 / file_in['imageWidth']
                y1, y2 = y1 / file_in['imageHeight'], y2 / file_in['imageHeight']
                cx, cy = (x1 + x2) / 2, (y1 + y2) / 2  # 中心点归一化的x坐标和y坐标
                wi, hi = abs(x2 - x1), abs(y2 - y1)  # 归一化的目标框宽度w，高度h
                line_content.append(str(cx))
                line_content.append(str(cy))
                line_content.append(str(wi))
                line_content.append(str(hi))
                # 使用空格连接列表中的所有元素，并写入文件
                file_handle.write(" ".join(line_content) + "\n")
 
        print("转换完成：", txt_filename)



def extract_matching_files(base_folder, target_folder, output_folder):
    """
    提取target_folder中与base_folder中同名(不含扩展名)的文件到到指定目录output_folder
    
    参数:
        base_folder (str): 包含.txt文档的文件夹路径
        images_folder (str): 包含.jpg图片的文件夹路径
        output_folder (str): 输出目录路径
    """

    os.makedirs(output_folder, exist_ok=True)
    
    # 获取文档文件名（不带扩展名）
    doc_files = set()
    for filename in os.listdir(base_folder):
        #if filename.endswith('.txt'):
            base_name = os.path.splitext(filename)[0]
            doc_files.add(base_name)
    
    # 查找匹配的图片并复制
    copied_count = 0
    for filename in os.listdir(target_folder):
        #if filename.endswith('.jpg') or filename.endswith('.jpeg'):
            base_name = os.path.splitext(filename)[0]
            if base_name in doc_files:
                src_path = os.path.join(target_folder, filename)
                dst_path = os.path.join(output_folder, filename)
                shutil.copy2(src_path, dst_path)
                copied_count += 1
                print(f"已复制: {filename}")
    
    print(f"共复制了 {copied_count} 文件到 {output_folder}")


def creat_folder_frame(folder_dir):
    '''
    创建yolo数据集文件夹框架
    '''
    images_dst = os.path.join(folder_dir, 'images')
    labels_dst = os.path.join(folder_dir, 'labels')
    with open(os.path.join(folder_dir,"data.yaml"), "w", encoding="utf-8") as f:
        f.write("path: \ntrain: \nval: \nnames: ")

    for folder in [images_dst, labels_dst]:
        for subfolder in ['train', 'val']:
            os.makedirs(os.path.join(folder, subfolder), exist_ok=True)
    


if __name__ == "__main__":
    #labelme2yolo_seg(["2point"],"labelme\label","labelme\yolo_label")
    labels = [
    "alarmIcon", "PEC-通信异常", "PSD-关", "PSD-未知", "PSD-禁止", "soundwave", "triangle", "wheelchair",
    "受电弓右-升弓", "受电弓右-故障", "受电弓右-降弓", "受电弓左-升弓", "受电弓左-故障", "受电弓左-降弓",
    "司机室激活-右侧", "司机室激活-左侧", "向前", "向后", "完全连接", "您的位置", "按钮-add", "按钮-close",
    "按钮-minus", "按钮-next", "按钮-previous", "按钮-record", "按钮-speaker_start", "按钮-speaker_stop",
    "按钮-喷油", "按钮-清除筛选", "按钮-设定", "按钮-返回", "提问警告", "操作提示", "机械连接", "灰",
    "烟火探测器-alarm", "烟火探测器-正常", "烟火探测器-污染", "烟火探测器-通信故障", "绿", "车门-close",
    "车门-inhibited", "车门-isolated", "车门-major_fault", "车门-minor_fault", "车门-obstacle_close",
    "车门-obstacle_open", "车门-open", "车门-urgency_unlock", "车门-通信异常", "黄"
    ]

    labelme2yolo_det(labels, "json", "txt")
    #extract_matching_images("yolo-dataset\\txt","yolo-dataset\imgs","yolo-dataset\img")
    #creat_folder_frame("yolo-dataset\dataset")

    # extract_matching_files("yolo-dataset\dataset\images\\train","yolo-dataset\\txt","yolo-dataset\dataset\labels\\train")
